import os
from utils.util import get_loader, get_optim, get_loss_fn, get_lr_scheduler
import torch
from tqdm import tqdm
from utils.setup import setup
from utils.log import get_loggere


class Model():
    def __init__(self, args):

        self.device = 'cpu'
        if torch.cuda.is_available():
            self.device = 'cuda'

        self.train_loader, self.test_loader = get_loader(args)
        print('data has been loaded over!')

        self.epochs = args.epoch
        # 模型保存路径
        self.save_model_path = args.save_model_path

        self.model_name = args.model
        self.model = setup(args).to(self.device)

        self.optimizer = get_optim(args, self.model)
        self.scheduler = get_lr_scheduler(args, self.optimizer)

        self.loss_fn = get_loss_fn(args)
        self.train_logger, self.test_logger = get_loggere(args)

    def train(self):
        num_batches = len(self.train_loader)
        self.model.train()
        with tqdm(total=num_batches) as pbar:
            for batch, (x, y) in enumerate(self.train_loader):
                x, y = x.to(self.device), y.to(self.device)
                # Compute prediction error
                pred = self.model(x)
                loss = self.loss_fn(pred, y)
                # Backpropagation
                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()
                # 添加日志信息
                self.train_logger.info(f' | {batch + 1}/{num_batches}]  training loss: {loss.item():0.6f}')
                pbar.set_postfix({'loss': f'{loss.item():0.6f}'})
                pbar.update(1)

            if self.scheduler:
                self.scheduler.step()
            loss = loss.item()
            tqdm.write(f"train loss: {loss:0.6f}")

    def test(self):
        # size = len(self.test_loader.dataset)
        num_batches = len(self.test_loader)
        self.model.eval()
        test_loss, all_acc = 0, 0
        with torch.no_grad():
            with tqdm(total=num_batches) as pbar:
                for batch, (x, y) in enumerate(self.test_loader):
                    x, y = x.to(self.device), y.to(self.device)
                    pred = self.model(x)
                    loss = self.loss_fn(pred, y).item()
                    acc = (pred.argmax(1) == y).sum().item() / len(y)
                    test_loss += loss
                    all_acc += acc
                    # 添加日志信息
                    self.test_logger.info(
                        f' | {batch + 1}/{num_batches}]    testing loss:{loss:0.6f}    testing acc:{acc:0.4f}')
                    pbar.update(1)

        test_loss /= num_batches
        all_acc /= num_batches
        tqdm.write(f"Test Accuracy: {all_acc:0.4f}, Avg loss: {test_loss:0.6f} \n")
        self.test_logger.info(
            f']    aver testing loss:{test_loss:0.6f}    aver testing acc:{all_acc:0.4f}')
        return all_acc

    def save_model(self):
        if not os.path.exists(self.save_model_path):
            os.mkdir(self.save_model_path)

        # 删除文件夹内所有模型，只保存最新的模型
        for filename in os.listdir(self.save_model_path):
            file_path = os.path.join(self.save_model_path, filename)
            try:
                if os.path.isfile(file_path) or os.path.islink(file_path):
                    os.unlink(file_path)
                elif os.path.isdir(file_path):
                    os.rmdir(file_path)
            except Exception as e:
                print(f'Failed to delete {file_path}. Reason: {e}')

        torch.save(self.model.state_dict(),
                   os.path.join(self.save_model_path, f'{self.model_name}_weights.pth'))

    def run(self):
        if not os.path.exists(self.save_model_path):
            os.mkdir(self.save_model_path)

        best_acc = 0
        for t in range(self.epochs):
            print(f"----------Epoch {t + 1} ---------")
            # 添加日志信息
            self.train_logger.info(f'] epoch {t + 1}/{self.epochs} begin')
            self.test_logger.info(f'] epoch {t + 1}/{self.epochs} begin')
            self.train()
            acc = self.test()
            if acc > best_acc:
                best_acc = acc
                self.save_model()
        print("Done!")
